import sys
import os
import argparse
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter
from io import BytesIO
import cv2
import traceback
from image_match import ImageFeature, MatchResult, SIFTMatcher

from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QGridLayout, QLabel, QHBoxLayout, QVBoxLayout,
    QFileDialog, QDialog, QDialogButtonBox, QTextEdit, QSizePolicy, QMenu, QAction
)
from PyQt5.QtCore import (
    Qt, QCoreApplication, QPoint, QPointF, QRect, QRectF, QSize, QSizeF, pyqtSignal,
    QTimer, QEasingCurve, QPropertyAnimation
)
from PyQt5.QtGui import (
    QPixmap, QImage, QPen, QPainter, QMouseEvent, QKeyEvent, QWheelEvent, QContextMenuEvent, QColor,
    QTransform, QFontMetrics, QDragEnterEvent, QDropEvent, QFontDatabase
)

import image_metric


class UsageDialog(QDialog):
    def __init__(self, usage_text: str, parent=None):
        super().__init__(parent)
        self.setWindowTitle("使用说明")
        self.resize(720, 520)
        layout = QVBoxLayout(self)
        text = QTextEdit(self)
        text.setReadOnly(True)
        mono = QFontDatabase.systemFont(QFontDatabase.FixedFont)
        text.setFont(mono)
        text.setPlainText(usage_text)
        text.setLineWrapMode(QTextEdit.WidgetWidth)
        layout.addWidget(text)
        buttons = QDialogButtonBox(QDialogButtonBox.Close, parent=self)
        buttons.rejected.connect(self.reject)
        buttons.accepted.connect(self.accept)
        layout.addWidget(buttons)


class ToastWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowFlags(
            Qt.ToolTip | Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint
        )
        self.setAttribute(Qt.WA_TranslucentBackground, True)

        self._label = QLabel(self)
        self._label.setStyleSheet("""
        QLabel {
            color: white;
            background-color: rgba(0, 0, 0, 180);
            border-radius: 6px;
            padding: 8px 12px;
            font-size: 12pt;
        }
        """)
        self._label.setAlignment(Qt.AlignCenter)

        self._timer = QTimer(self)
        self._timer.setSingleShot(True)
        self._timer.timeout.connect(self._on_timeout)

        self._fade_anim = QPropertyAnimation(self, b"windowOpacity", self)
        self._fade_anim.setEasingCurve(QEasingCurve.InOutQuad)

        self._duration_ms = 500

    def _safe_disconnect_finished(self):
        try:
            self._fade_anim.finished.disconnect(self.hide)
        except Exception:
            pass
        try:
            self._fade_anim.finished.disconnect(self._on_fade_out_finished)
        except Exception:
            pass

    def show_message(self, message: str, duration_ms: int = 500):
        self._timer.stop()
        self._fade_anim.stop()
        self._safe_disconnect_finished()

        self._duration_ms = duration_ms
        self._label.setText(message)
        self._label.adjustSize()

        margin = 12
        w = self._label.width() + margin
        h = self._label.height() + margin
        self.resize(w, h)
        self._label.move((w - self._label.width()) // 2,
                         (h - self._label.height()) // 2)

        if self.parent() is not None:
            pw = self.parent().width()
            ph = self.parent().height()
            local_pos = QPoint((pw - w) // 2, int(ph * 0.85) - h // 2)
            global_pos = self.parent().mapToGlobal(local_pos)
            x, y = global_pos.x(), global_pos.y()
        else:
            screen = QApplication.primaryScreen().availableGeometry()
            x = screen.x() + (screen.width() - w) // 2
            y = screen.y() + int(screen.height() * 0.85) - h // 2

        self.move(x, y)

        self.setWindowOpacity(0.0)
        self.show()
        self._fade_anim.stop()
        self._fade_anim.setDuration(120)
        self._fade_anim.setStartValue(0.0)
        self._fade_anim.setEndValue(1.0)
        self._fade_anim.start()

        self._timer.start(self._duration_ms)

    def _on_timeout(self):
        self._fade_anim.stop()
        self._safe_disconnect_finished()
        self._fade_anim.setDuration(180)
        self._fade_anim.setStartValue(1.0)
        self._fade_anim.setEndValue(0.0)
        self._fade_anim.finished.connect(self._on_fade_out_finished)
        self._fade_anim.start()

    def _on_fade_out_finished(self):
        self._safe_disconnect_finished()
        self.hide()


class CanvasWidget(QWidget):
    offset_scale_changed = pyqtSignal(QPointF, float)
    position_changed = pyqtSignal(QPointF)
    rectangle_changed = pyqtSignal(QRectF)
    delete_requested = pyqtSignal()

    def __init__(self, parent_viewer):
        super().__init__(parent_viewer)
        self.viewer = parent_viewer
        self.setMouseTracking(True)
        self.show_context_menu = False

    def sizeHint(self):
        return QSize(200, 200)

    def paintEvent(self, event):
        if self.viewer.pixmap is None:
            return
        painter = QPainter(self)
        scaled_size = QSizeF(self.viewer.pixmap.size()) * self.viewer.scale
        img_rect = QRectF(self.viewer.offset, scaled_size)
        visible_rect = QRectF(0, 0, self.width(), self.height())
        draw_rect = img_rect.intersected(visible_rect)
        if not draw_rect.isEmpty():
            source_rect = QRectF(
                (draw_rect.x() - self.viewer.offset.x()) / self.viewer.scale,
                (draw_rect.y() - self.viewer.offset.y()) / self.viewer.scale,
                draw_rect.width() / self.viewer.scale,
                draw_rect.height() / self.viewer.scale
            )
            painter.drawPixmap(draw_rect, self.viewer.play_pixmap, source_rect)

    def mousePressEvent(self, event: QMouseEvent):
        self.viewer.press_start_pos = event.pos()
        self.viewer.move_pos = event.pos()
            

    def mouseMoveEvent(self, event: QMouseEvent):
        if event.buttons() & Qt.RightButton:
            delta = event.pos() - self.viewer.move_pos
            self.viewer.offset += delta
            self.viewer.move_pos = event.pos()
            self.update()
            self.offset_scale_changed.emit(self.viewer.offset, self.viewer.scale)
        elif event.buttons() & Qt.LeftButton:
            pos = event.pos()
            delta = pos - self.viewer.press_start_pos
            dis = abs(delta.x()) + abs(delta.y())
            if dis > 3:
                img_start_x = (self.viewer.press_start_pos.x() - self.viewer.offset.x()) / self.viewer.scale
                img_start_y = (self.viewer.press_start_pos.y() - self.viewer.offset.y()) / self.viewer.scale
                img_end_x = (pos.x() - self.viewer.offset.x()) / self.viewer.scale
                img_end_y = (pos.y() - self.viewer.offset.y()) / self.viewer.scale
                rect = QRectF(img_start_x, img_start_y, img_end_x - img_start_x, img_end_y - img_start_y)
                self.viewer.updateRectangle(rect)
                self.update()

    def mouseReleaseEvent(self, event: QMouseEvent):
        pos = event.pos()
        delta = pos - self.viewer.press_start_pos
        dis = abs(delta.x()) + abs(delta.y())
        if event.button() == Qt.LeftButton:
            if dis < 3:
                img_x = (self.viewer.press_start_pos.x() - self.viewer.offset.x()) / self.viewer.scale
                img_y = (self.viewer.press_start_pos.y() - self.viewer.offset.y()) / self.viewer.scale
                if 0 <= img_x < self.viewer.image.width() and 0 <= img_y < self.viewer.image.height():
                    self.position_changed.emit(QPointF(img_x, img_y))
            else:
                img_start_x = (self.viewer.press_start_pos.x() - self.viewer.offset.x()) / self.viewer.scale
                img_start_y = (self.viewer.press_start_pos.y() - self.viewer.offset.y()) / self.viewer.scale
                img_end_x = (pos.x() - self.viewer.offset.x()) / self.viewer.scale
                img_end_y = (pos.y() - self.viewer.offset.y()) / self.viewer.scale
                if 0 <= img_start_x < self.viewer.image.width() and 0 <= img_start_y < self.viewer.image.height() and \
                0 <= img_end_x < self.viewer.image.width() and 0 <= img_end_y < self.viewer.image.height():
                    left = min(img_start_x, img_end_x)
                    top = min(img_start_y, img_end_y)
                    right = max(img_start_x, img_end_x)
                    bottom = max(img_start_y, img_end_y)
                    rect = QRectF(QPointF(left, top), QPointF(right, bottom))
                    self.rectangle_changed.emit(rect)
        elif event.button() == Qt.RightButton:
            self.show_context_menu = dis < 3
                

    def wheelEvent(self, event: QWheelEvent):
        delta = event.angleDelta().y()
        old_center = QPointF(
            (self.width() / 2.0 - self.viewer.offset.x()) / self.viewer.scale,
            (self.height() / 2.0 - self.viewer.offset.y()) / self.viewer.scale
        )
        if delta > 0:
            self.viewer.scale = min(self.viewer.scale * 1.1, self.viewer.max_scale)
        else:
            self.viewer.scale = max(self.viewer.scale * 0.9, self.viewer.min_scale)
        self.viewer.offset = QPointF(
            self.width() / 2.0 - old_center.x() * self.viewer.scale,
            self.height() / 2.0 - old_center.y() * self.viewer.scale
        )
        self.offset_scale_changed.emit(self.viewer.offset, self.viewer.scale)
        self.update()

    def contextMenuEvent(self, event: QContextMenuEvent):
        # 检查是否是真正的右键点击（而不是拖拽）
        # 如果右键拖拽距离大于阈值，则不显示菜单
        if not self.show_context_menu:
            return
            
        # 创建右键菜单
        menu = QMenu(self)
        
        # 添加删除动作
        delete_action = QAction("删除当前图片", self)
        delete_action.triggered.connect(self.delete_requested.emit)
        menu.addAction(delete_action)
        
        # 显示菜单
        menu.exec_(event.globalPos())


class ImageViewer(QWidget):
    offset_scale_changed = pyqtSignal(QPointF, float)
    position_changed = pyqtSignal(QPointF)
    rectangle_changed = pyqtSignal(QRectF)
    delete_requested = pyqtSignal()

    def __init__(self, image_path, name):
        super().__init__()
        self.image_path = image_path
        self.name = name
        self.is_show_normalization = False
        self.color_space = 'yuv'
        self.rotation_angle = 0

        self.scale = 1.0
        self.offset = QPointF(0, 0)
        self.min_scale = 0.1
        self.max_scale = 40.0
        self.press_start_pos = QPointF(0, 0)
        self.move_pos = QPointF(0, 0)
        self.select_pos = None
        self.select_rect = None
        self.matcher = SIFTMatcher()
        self.original_image = QImage(self.image_path)
        if self.original_image.isNull():
            print(f"Error: Unable to load image {self.image_path}.")
            sys.exit(1)

        self.image = self.original_image.copy()
        self.pixmap = QPixmap.fromImage(self.image)
        self.play_pixmap = self.pixmap.copy()

        self._build_ui()

    def _build_ui(self):
        self.setMinimumSize(200, 200)

        # 中部绘制区域（Canvas）占满
        self.canvas = CanvasWidget(self)
        self.canvas.offset_scale_changed.connect(self._relay_offset_scale)
        self.canvas.position_changed.connect(self._relay_position)
        self.canvas.rectangle_changed.connect(self._relay_rectangle)
        self.canvas.delete_requested.connect(self.delete_viewer)

        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(0, 0, 0, 0)
        main_layout.setSpacing(0)
        main_layout.addWidget(self.canvas, 1)
        self.setLayout(main_layout)

        # 顶部半透明 info_bar 作为浮层（不占据布局空间）
        self.info_bar = QWidget(self)
        self.info_bar.setStyleSheet("background-color: rgba(0, 0, 0, 100);")

        info_layout = QHBoxLayout()
        info_layout.setContentsMargins(0, 0, 0, 0)
        info_layout.setSpacing(0)

        self.name_label = QLabel(self.name, self.info_bar)
        self.name_label.setObjectName("name_label")
        self.name_label.setToolTip(self.image_path)


        self.position_color_label = QLabel("Position: (0, 0), YUV:(0,0,0)", self.info_bar)
        self.position_color_label.setObjectName("position_color_label")
        # 设置悬浮提示：显示完整文本
        self.position_color_label.setToolTip("Position: (0, 0), YUV:(0,0,0)")

        self.setStyleSheet("""QToolTip {
            color: #ffeb3b;
            background-color: rgba(0,0,0,200);
            border: 1px solid #888;
            padding: 4px;
            font-size: 12px;
            font-family: "Microsoft YaHei", "Segoe UI", Arial;
        }
        QLabel#name_label,
        QLabel#position_color_label{
            color: #ffeb3b;
            background-color: transparent;
            border: 1px solid #888;
            padding: 3px;
            font-size: 12px;
            font-family: "Microsoft YaHei", "Segoe UI", Arial;
        }
        """)

        info_layout.addWidget(self.name_label)
        info_layout.addStretch()
        info_layout.addWidget(self.position_color_label)
        self.info_bar.setLayout(info_layout)
        self.info_bar.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Minimum)
        # 初始化放置到顶部
        self.info_bar.move(0, 0)
        self.info_bar.raise_()  # 置顶覆盖
        # 右下角直方图
        self.hist_label = QLabel(self)
        self.hist_label.setStyleSheet("background-color: rgba(255, 255, 255, 200);")
        self.hist_label.setAlignment(Qt.AlignCenter)
        self.hist_label.hide()

        # 初次设置文本省略
        self._set_name_label_text()
        # 初次根据字体计算并设置 info_bar 高度与位置
        self.update_info_bar_height()

    def update_info_bar_height(self):
        h1 = self.name_label.sizeHint().height()
        h2 = self.position_color_label.sizeHint().height()
        text_h = max(h1, h2)
        bar_h = max(30, text_h + 3)
        self.info_bar.resize(self.width(), bar_h)
        self.info_bar.move(0, 0)
        self.info_bar.raise_()

    def _relay_offset_scale(self, offset, scale):
        self.offset_scale_changed.emit(offset, scale)

    def _relay_position(self, pos):
        self.position_changed.emit(pos)

    def _relay_rectangle(self, rect):
        self.rectangle_changed.emit(rect)

    def _shorten_path_for_label(self, full_path, keep_levels=2):
        dir_path = os.path.dirname(full_path)
        norm = dir_path.replace('\\', '/')
        parts = [p for p in norm.split('/') if p]
        if len(parts) == 0:
            return ""
        if len(parts) <= keep_levels:
            short = dir_path
        else:
            short = '.../' + '/'.join(parts[-keep_levels:])
        return short

    def _set_name_label_text(self, max_width_px=None):
        base = os.path.basename(self.image_path) if self.image_path else self.name
        short_parent = self._shorten_path_for_label(self.image_path, keep_levels=2)
        show_text = base if not short_parent else f"{base}  ({short_parent})"

        label_width = max_width_px if max_width_px is not None else max(120, int(self.width() * 0.45))
        fm = QFontMetrics(self.name_label.font())
        elided = fm.elidedText(show_text, Qt.ElideMiddle, label_width)
        self.name_label.setText(elided)
        self.name_label.setToolTip(self.image_path)

    def resizeEvent(self, event):
        self.update_info_bar_height()
        available = max(120, int(self.width() * 0.45))
        self._set_name_label_text(max_width_px=available)

        # 控制直方图不超过 viewer 尺寸
        if self.hist_label.isVisible():
            margin = 5
            max_w = max(0, self.width() - margin * 2)
            max_h = max(0, self.height() - margin * 2)
            if max_w <= 0 or max_h <= 0:
                self.hist_label.hide()
            else:
                # 当前 pixmap 按需要再次限制（只缩小不放大）
                pm = self.hist_label.pixmap()
                if pm is not None and not pm.isNull():
                    pw, ph = pm.width(), pm.height()
                    new_w, new_h = pw, ph
                    if pw > max_w or ph > max_h:
                        scale = min(max_w / pw, max_h / ph)
                        new_w = max(1, int(pw * scale))
                        new_h = max(1, int(ph * scale))
                        self.hist_label.resize(new_w, new_h)
                    else:
                        # 若已有尺寸大于 pixmap 本身，收缩为 pixmap 尺寸
                        self.hist_label.resize(pw, ph)
                    self.hist_label.move(self.width() - self.hist_label.width() - margin,
                                        self.height() - self.hist_label.height() - margin)

        super().resizeEvent(event)


    def updateGeometry(self, offset: QPointF, scale: float):
        self.offset = offset
        self.scale = scale
        self.canvas.update()

    def rotateImage(self, clockwise=True):
        if clockwise:
            self.rotation_angle = (self.rotation_angle + 90) % 360
        else:
            self.rotation_angle = (self.rotation_angle - 90) % 360

        transform = QTransform()
        transform.rotate(self.rotation_angle)
        self.image = self.original_image.transformed(transform, Qt.SmoothTransformation)
        self.pixmap = QPixmap.fromImage(self.image)
        self.play_pixmap = self.pixmap.copy()

        self.select_pos = None
        self.select_rect = None
        self.hist_label.hide()
        # 同步 tooltip
        self._set_position_color_text("Position: (0, 0), YUV:(0,0,0)")
        self.offset = QPointF(0, 0)
        self.scale = 1.0
        self.canvas.update()

    def rgbToHsv(self, rgb):
        rgb = np.asarray(rgb, dtype=float).reshape(-1)
        if rgb.size != 3:
            raise ValueError("rgb must be an iterable of length 3, e.g. (R, G, B).")
        r, g, b = rgb
        if r > 1.0 or g > 1.0 or b > 1.0:
            r, g, b = r / 255.0, g / 255.0, b / 255.0
        cmax = max(r, g, b)
        cmin = min(r, g, b)
        delta = cmax - cmin
        eps = 1e-10
        if delta < eps:
            h = 0.0
        elif cmax == r:
            h = 60.0 * (((g - b) / (delta + eps)) % 6.0)
        elif cmax == g:
            h = 60.0 * (((b - r) / (delta + eps)) + 2.0)
        else:
            h = 60.0 * (((r - g) / (delta + eps)) + 4.0)
        if cmax < eps:
            s = 0.0
        else:
            s = delta / (cmax + eps)
        v = cmax
        h = h % 360.0
        return np.array([h, s, v], dtype=float)

    def rgbToYuv(self, rgb):
        rgb = rgb / 255.0
        transform_matrix = np.array([
            [0.299, 0.587, 0.114],
            [-0.14713, -0.28886, 0.436],
            [0.615, -0.51499, -0.10001]
        ])
        yuv = np.dot(transform_matrix, rgb)
        yuv[1] += 0.5
        yuv[2] += 0.5
        yuv = np.clip(yuv * 255, 0, 255).astype(np.uint8)
        return yuv

    def _set_position_color_text(self, full_text: str):
        # 统一更新 position_color_label 的文本与 tooltip
        self.position_color_label.setText(full_text)
        self.position_color_label.setToolTip(full_text)

    def updatePosition(self, pos: QPointF):
        self.select_pos = pos
        self.select_rect = None
        self.hist_label.hide()
        x, y = int(pos.x()), int(pos.y())
        if 0 <= x < self.image.width() and 0 <= y < self.image.height():
            color = self.image.pixelColor(x, y)
            rgba = np.array([color.red(), color.green(), color.blue(), color.alpha()])
            if self.color_space == 'yuv':
                yuv = self.rgbToYuv(rgba[:3])
                if self.is_show_normalization:
                    yuv = yuv / 255.0
                    text = f"Position: ({x}, {y}), YUV:({yuv[0]:0.2f}, {yuv[1]:0.2f}, {yuv[2]:0.2f})"
                else:
                    text = f"Position: ({x}, {y}), YUV:({int(yuv[0])}, {int(yuv[1])}, {int(yuv[2])})"
            elif self.color_space == 'rgba':
                if self.is_show_normalization:
                    rgba = rgba / 255.0
                    text = f"Position: ({x}, {y}), RGBA:({rgba[0]:0.2f}, {rgba[1]:0.2f}, {rgba[2]:0.2f}, {rgba[3]:0.2f})"
                else:
                    text = f"Position: ({x}, {y}), RGBA:({rgba[0]}, {rgba[1]}, {rgba[2]}, {rgba[3]})"
            else:
                hsv_val = self.rgbToHsv(rgba[:3].astype(np.float32))
                text = f"Position: ({x}, {y}), HSV:({hsv_val[0]:.1f}, {hsv_val[1]:.3f}, {hsv_val[2]:.3f})"
            self._set_position_color_text(text)
        else:
            self._set_position_color_text("Position: (0, 0), RGBA: (0, 0, 0, 0)")

        self.play_pixmap = self.pixmap.copy()
        line_lenght = 10
        line_width = 1
        painter = QPainter(self.play_pixmap)
        pen = QPen(QColor(255, 0, 0, 100), line_width)
        painter.setPen(pen)
        painter.drawLine(x - line_lenght, y, x - 1, y)
        painter.drawLine(x + 1, y, x + line_lenght, y)
        painter.drawLine(x, y - line_lenght, x, y - 1)
        painter.drawLine(x, y + 1, x, y + line_lenght)
        painter.end()
        self.canvas.update()

    def calYUVRange(self, image):
        yuv = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
        y_range = (np.min(yuv[:, :, 0]), np.max(yuv[:, :, 0]))
        u_range = (np.min(yuv[:, :, 1]), np.max(yuv[:, :, 1]))
        v_range = (np.min(yuv[:, :, 2]), np.max(yuv[:, :, 2]))
        return y_range, u_range, v_range

    def calMultiChannelHist(self, image, channels, channel_names, output_width, output_height, dpi=100):
        fig_width = output_width / dpi
        fig_height = output_height / dpi
        plt.figure(figsize=(fig_width, fig_height), dpi=dpi)
        
        for i, channel in enumerate(channels):
            ax = plt.subplot(1, len(channels), i + 1)
            channel_data = image[:, :, channel]
            
            # 检查数据有效性
            if channel_data.size == 0:
                # 如果没有数据，显示空图表
                plt.title(f"{channel_names[i]}(N/A)")
                plt.xlim(0, 255)
                plt.ylim(0, 1)
                continue
                
            channel_mean = float(np.mean(channel_data))
            data_min = channel_data.min()
            data_max = channel_data.max()
            
            # 改进的数据范围处理
            if data_min == data_max:
                # 如果所有值相同，创建一个合理的范围
                if data_min == 0:
                    data_min, data_max = -0.5, 0.5
                else:
                    # 根据数据类型调整范围
                    if np.issubdtype(channel_data.dtype, np.integer):
                        data_min = data_min - 0.5
                        data_max = data_max + 0.5
                    else:
                        range_size = max(0.1, abs(data_min) * 0.1)
                        data_min = data_min - range_size
                        data_max = data_max + range_size
            else:
                # 确保最小范围，避免数据过于集中
                data_range = data_max - data_min
                if data_range < 1e-6:
                    data_min = data_min - 0.5
                    data_max = data_max + 0.5
            
            # 计算频数直方图（density=False）
            hist, bins = np.histogram(channel_data, bins=20, density=False, range=(data_min, data_max))
            
            # 确保直方图值非负
            hist = np.maximum(hist, 0)
            
            # 绘制直方图
            plt.bar(bins[:-1], hist, width=np.diff(bins), align="edge", color='skyblue', edgecolor='black')
            
            # 设置x轴ticks
            ticks = np.linspace(data_min, data_max, num=5, dtype=np.float32)
            plt.xticks(ticks, rotation=-45)
            
            # 格式化x轴标签
            if np.issubdtype(channel_data.dtype, np.integer):
                plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%d'))
            else:
                if data_max - data_min < 0.5:
                    plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%.2f'))
                else:
                    plt.gca().xaxis.set_major_formatter(FormatStrFormatter('%.1f'))
            
            # 设置y轴为整数刻度，显示实际频数
            max_hist = hist.max() if hist.size > 0 else 1
            plt.ylim(0, max_hist * 1.1)  # 留10%的余量
            
            # 设置y轴刻度为整数
            y_ticks = np.arange(0, int(max_hist * 1.1) + 1, max(1, int(max_hist * 1.1) // 5))
            plt.yticks(y_ticks)
            
            # 添加标题
            plt.title(f"{channel_names[i]}({channel_mean:.1f})")
            
            # 添加网格
            plt.grid(axis='y', alpha=0.5)

        plt.tight_layout()
        buf = BytesIO()
        plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=dpi)
        buf.seek(0)
        hist_img = np.frombuffer(buf.getvalue(), dtype=np.uint8)
        hist_img = cv2.imdecode(hist_img, cv2.IMREAD_COLOR)
        plt.close()
        hist_img = cv2.resize(hist_img, (output_width, output_height))
        return hist_img

    def updateRectangle(self, rect: QRectF, statistics=False):
        self.select_rect = rect
        self.select_pos = None

        if statistics:
            rgb = self.qimage_to_numpy(rect.toRect())
            if rgb is not None:
                if self.color_space == 'yuv':
                    yuv = cv2.cvtColor(rgb, cv2.COLOR_RGB2YUV)
                    if self.is_show_normalization:
                        yuv = yuv.astype(np.float32) / 255.0
                    y_range = (np.min(yuv[:, :, 0]), np.max(yuv[:, :, 0]))
                    u_range = (np.min(yuv[:, :, 1]), np.max(yuv[:, :, 1]))
                    v_range = (np.min(yuv[:, :, 2]), np.max(yuv[:, :, 2]))
                    y_std = np.std(yuv[:, :, 0])
                    u_std = np.std(yuv[:, :, 1])
                    v_std = np.std(yuv[:, :, 2])
                    if self.is_show_normalization:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"yuv: ({y_range[0]:.2f}-{y_range[1]:.2f}, {u_range[0]:.2f}-{u_range[1]:.2f}, {v_range[0]:.2f}-{v_range[1]:.2f}), "
                            f"std: ({y_std:.2f}, {u_std:.2f}, {v_std:.2f})"
                        )
                    else:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"yuv: ({y_range[0]}-{y_range[1]}, {u_range[0]}-{u_range[1]}, {v_range[0]}-{v_range[1]}), "
                            f"std: ({y_std:.2f}, {u_std:.2f}, {v_std:.2f})"
                        )
                elif self.color_space == 'rgba':
                    rgba = cv2.cvtColor(rgb, cv2.COLOR_BGRA2RGBA)
                    if self.is_show_normalization:
                        rgba = rgba.astype(np.float32) / 255.0
                    r_range = (np.min(rgba[:, :, 0]), np.max(rgba[:, :, 0]))
                    g_range = (np.min(rgba[:, :, 1]), np.max(rgba[:, :, 1]))
                    b_range = (np.min(rgba[:, :, 2]), np.max(rgba[:, :, 2]))
                    a_range = (np.min(rgba[:, :, 3]), np.max(rgba[:, :, 3]))
                    r_std = np.std(rgba[:, :, 0])
                    g_std = np.std(rgba[:, :, 1])
                    b_std = np.std(rgba[:, :, 2])
                    a_std = np.std(rgba[:, :, 3])
                    if self.is_show_normalization:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"rgba: ({r_range[0]:.2f}-{r_range[1]:.2f}, {g_range[0]:.2f}-{g_range[1]:.2f}, {b_range[0]:.2f}-{b_range[1]:.2f}, {a_range[0]:.2f}-{a_range[1]:.2f}), "
                            f"std: ({r_std:.2f}, {g_std:.2f}, {b_std:.2f}, {a_std:.2f})"
                        )
                    else:
                        info_text = (
                            f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                            f"rgba: ({r_range[0]}-{r_range[1]}, {g_range[0]}-{g_range[1]}, {b_range[0]}-{b_range[1]}, {a_range[0]}-{a_range[1]}), "
                            f"std: ({r_std:.2f}, {g_std:.2f}, {b_std:.2f}, {a_std:.2f})"
                        )
                elif self.color_space == 'hsv':
                    rgb = rgb.astype(np.float32) / 255.0
                    hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
                    h_range = (np.min(hsv[:, :, 0]), np.max(hsv[:, :, 0]))
                    s_range = (np.min(hsv[:, :, 1]), np.max(hsv[:, :, 1]))
                    v_range = (np.min(hsv[:, :, 2]), np.max(hsv[:, :, 2]))
                    h_std = np.std(hsv[:, :, 0])
                    s_std = np.std(hsv[:, :, 1])
                    v_std = np.std(hsv[:, :, 2])
                    info_text = (
                        f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                        f"hsv: ({h_range[0]:.1f}-{h_range[1]:.1f}, {s_range[0]:.2f}-{s_range[1]:.2f}, {v_range[0]:.2f}-{v_range[1]:.2f}), "
                        f"std: ({h_std:.2f}, {s_std:.2f}, {v_std:.2f})"
                    )
                # 更新文本与 tooltip
                self._set_position_color_text(info_text)

        self.play_pixmap = self.pixmap.copy()
        line_width = 1
        painter = QPainter(self.play_pixmap)
        pen = QPen(QColor(255, 0, 0, 100), line_width)
        painter.setPen(pen)
        painter.drawRect(rect)
        painter.end()
        self.canvas.update()

    def updateHist(self, rect: QRectF):
        rgb = self.qimage_to_numpy(rect.toRect())
        if rgb is not None:
            if self.color_space == 'yuv':
                yuv = cv2.cvtColor(rgb, cv2.COLOR_RGB2YUV)
                if self.is_show_normalization:
                    yuv = yuv.astype(np.float32) / 255.0
                hist_img = self.calMultiChannelHist(
                    yuv, [0, 1, 2], ['Y', 'U', 'V'], output_width=int(self.width()-6), output_height=int(self.height()/3)
                )
            elif self.color_space == 'rgba':
                if self.is_show_normalization:
                    rgb = rgb.astype(np.float32) / 255.0
                hist_img = self.calMultiChannelHist(
                    rgb[:, :, :3], [0, 1, 2], ['R', 'G', 'B'], output_width=int(self.width()-6), output_height=int(self.height()/3)
                )
            else:
                rgb = rgb.astype(np.float32) / 255.0
                hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
                hist_img = self.calMultiChannelHist(
                    hsv, [0, 1, 2], ['H', 'S', 'V'], output_width=int(self.width()-6), output_height=int(self.height()/3)
                )

            # 限制 hist_label 不超过 viewer 尺寸
            margin = 5
            max_w = max(0, self.width() - margin * 2)
            max_h = max(0, self.height() - margin * 2)

            if max_w == 0 or max_h == 0:
                self.hist_label.hide()
                return

            h, w, _ = hist_img.shape
            scale = min(max_w / w, max_h / h, 1.0)  # 只缩小不放大
            new_w = max(1, int(w * scale))
            new_h = max(1, int(h * scale))
            if (new_w != w) or (new_h != h):
                hist_img = cv2.resize(hist_img, (new_w, new_h), interpolation=cv2.INTER_AREA)

            height, width, channel = hist_img.shape
            bytes_per_line = 3 * width
            q_hist_img = QImage(hist_img.data, width, height, bytes_per_line, QImage.Format_RGB888)
            self.hist_label.setPixmap(QPixmap.fromImage(q_hist_img))
            self.hist_label.resize(width, height)
            self.hist_label.move(self.width() - width - margin, self.height() - height - margin)
            self.hist_label.show()


    def updateMTF(self, rect: QRectF):
        rgb = self.qimage_to_numpy(rect.toRect())
        if rgb is not None:
            yuv = cv2.cvtColor(rgb, cv2.COLOR_RGB2YUV)
            luma = yuv[:, :, 0]
            mtf_50, nyquist = image_metric.compute_mtf(luma)
            if mtf_50 >= 0 and nyquist >= 0:
                text = (
                    f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), "
                    f"MTF50:{mtf_50:.2f}, Nyquist:{nyquist:.2f}%"
                )
            else:
                text = (
                    f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), MTF计算失败"
                )
            self._set_position_color_text(text)

    def updatePSNR(self, rect: QRectF):
        rgb = self.qimage_to_numpy(rect.toRect())
        if rgb is not None:
            if self.color_space == 'yuv':
                yuv = cv2.cvtColor(rgb, cv2.COLOR_RGB2YUV)
                if self.is_show_normalization:
                    yuv = yuv.astype(np.float32) / 255.0
                src_img = yuv
            elif self.color_space == 'rgba':
                if self.is_show_normalization:
                    rgb = rgb.astype(np.float32) / 255.0
                src_img = rgb
            else:
                rgb = rgb.astype(np.float32) / 255.0
                hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
                src_img = hsv
            psnrs = image_metric.compute_psnr(src_img)
            psnrs_str = ', '.join([f"{psnr:.2f}" for psnr in psnrs])
            self._set_position_color_text(
                f"rect: ({int(rect.x())}, {int(rect.y())} | {int(rect.width())}, {int(rect.height())}), psnr(dB): ({psnrs_str})"
            )

    def saveCorp(self, rect: QRectF, saveDir):
        rgb = self.qimage_to_numpy(rect.toRect())
        if rgb is not None:
            cv2.imwrite(os.path.join(saveDir, os.path.basename(self.image_path)), rgb)

    def showPixelNormalization(self):
        self.is_show_normalization = not self.is_show_normalization
        if self.select_pos is not None:
            self.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            self.updateRectangle(self.select_rect, statistics=True)

    def updateColorSpace(self, color_space):
        self.color_space = color_space
        if self.select_pos is not None:
            self.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            self.updateRectangle(self.select_rect, statistics=True)

    def resetImage(self):
        self.image = self.original_image.copy()
        self.pixmap = QPixmap.fromImage(self.image)
        self.play_pixmap = self.pixmap.copy()
        self.canvas.update()

    def cropSubImage(self, rect: QRectF):
        """获取当前选区的ROI截图、SIFT特征点和偏移量
        
        Args:
            rect: 选区矩形
            
        Returns:
            tuple: (roi_image, roi_features, roi_offset)
                roi_image: ROI区域的QImage
                roi_features: ROI区域的SIFT特征点 (keypoints, descriptors)
                roi_offset: ROI区域的偏移量 (x, y)
        """
        if rect is None or rect.isEmpty():
            return None, None, None
            
        # 恢复image
        self.resetImage()
        # 获取ROI图像
        image = self.qimage_to_numpy()
        rint = rect.toRect()
        roi_rect = (rint.x(), rint.y(), rint.width(), rint.height())
        # 提取sift特征
        result = self.matcher.feature_extract(img = image, rect = roi_rect)
        return result


    def qimage_to_numpy(self, rect: QRect = None, fmt = 'rgb'):
        if rect is None:
            qimage = self.image.copy(self.image.rect()).convertToFormat(QImage.Format_ARGB32)
        else:
            qimage = self.image.copy(rect).convertToFormat(QImage.Format_ARGB32)
        if qimage.isNull():
            return None
        ptr = qimage.bits()
        ptr.setsize(qimage.byteCount())
        arr = np.frombuffer(ptr, dtype=np.uint8).reshape(qimage.height(), qimage.width(), 4)
        rgba = cv2.cvtColor(arr, cv2.COLOR_BGRA2RGBA)
        # 修正：正确处理ARGB32到不同颜色空间的转换
        if fmt == 'rgb':
            # ARGB32 -> BGRA -> RGB (正确顺序)
            dst = cv2.cvtColor(rgba, cv2.COLOR_RGBA2RGB)
        elif fmt == 'bgr':
            # ARGB32 -> BGR (直接转换，OpenCV默认格式)
            dst = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGR)
        elif fmt == 'rgba':
            # ARGB32 -> BGRA (保持原样)
            dst = arr.copy()
        elif fmt == 'bgra':
            # ARGB32 -> BGRA (保持原样)
            dst = cv2.cvtColor(arr, cv2.COLOR_RGBA2BGRA)
        return dst

    def delete_viewer(self):
        """删除当前ImageViewer"""
        # 发射删除信号，让父窗口处理删除逻辑
        self.delete_requested.emit()


    def alignment(self, feature:ImageFeature, target_offset, target_rect):
        """将当前图像与目标ROI区域进行对齐
        
        Args:
            feature: 目标ROI区域的SIFT特征点 (keypoints, descriptors)
            target_offset: 目标ROI区域的偏移量 (x, y)
            target_rect: 目标ROI区域的矩形
            
        Returns:
            bool: 对齐是否成功
        """
        if feature is None:
            return False
            
        try:
            # 恢复image
            self.resetImage()
            # 获取当前图像的SIFT特征 - 使用BGR格式确保OpenCV兼容性
            current_bgr = self.qimage_to_numpy()
            if current_bgr is None:
                return False
                
            
            # 特征匹配 - 使用BGR格式
            match_result = self.matcher.feature_match(feature=feature, img=current_bgr)
            if match_result is None:
                return False
                
            aligned_image = match_result.warped_img
            if aligned_image is None:
                return False
                
            # 检查对齐后图像的通道顺序
            
            # 确保对齐后的图像转换为RGB格式用于显示
            # 如果aligned_image是BGR格式（OpenCV默认），需要转换为RGB
            if len(aligned_image.shape) == 3 and aligned_image.shape[2] == 3:
                # 检查是否需要转换BGR到RGB
                try:
                    # 验证输入图像的颜色通道顺序
                    w = aligned_image.shape[1]
                    h = aligned_image.shape[0]
                    channels = aligned_image.shape[2]
                    stride = channels * w
                    self.image = QImage(aligned_image.data, w, h, stride, QImage.Format_RGB888)
                    
                except Exception as e:
                    # 如果转换失败，直接使用原始数据
                    w = aligned_image.shape[1]
                    h = aligned_image.shape[0]
                    channels = aligned_image.shape[2]
                    stride = channels * w
                    self.image = QImage(aligned_image.data, w, h, stride, QImage.Format_RGB888)
            else:
                # 灰度图像或其他格式
                w = aligned_image.shape[1]
                h = aligned_image.shape[0]
                if len(aligned_image.shape) == 3:
                    channels = aligned_image.shape[2]
                    stride = channels * w
                    self.image = QImage(aligned_image.data, w, h, stride, QImage.Format_RGB888)
                else:
                    # 灰度图像，需要转换为RGB
                    rgb_image = cv2.cvtColor(aligned_image, cv2.COLOR_GRAY2RGB)
                    w = rgb_image.shape[1]
                    h = rgb_image.shape[0]
                    channels = rgb_image.shape[2]
                    stride = channels * w
                    self.image = QImage(rgb_image.data, w, h, stride, QImage.Format_RGB888)
                    
            
            self.pixmap = QPixmap.fromImage(self.image)
            self.play_pixmap = self.pixmap.copy()
            
            # 同步offset
            self.offset = QPointF(target_offset)
            
            # 清除之前的选区
            self.select_pos = None
            self.select_rect = QRectF(target_rect)
            self.hist_label.hide()
            
            # 更新显示
            self.canvas.update()
            return True
        except Exception as e:
            traceback.print_exc()
            return False


class ImageComparator(QMainWindow):
    def __init__(self, image_paths):
        super().__init__()
        image_paths.sort()
        self.image_paths = image_paths[:9]
        self.color_space = 'yuv'
        self.viewers = []
        self.grid_layout = None
        self.container = None
        self.initUI()
        self.select_pos = None
        self.select_rect: QRectF = None
        self._active_viewer = None  # 记录当前操作的viewer

        self.setAcceptDrops(True)

    def get_grid_layout(self, count):
        if count == 1:
            return (1, 1)
        elif count == 2:
            return (1, 2)
        elif count == 3:
            return (2, 2)
        elif count == 4:
            return (2, 2)
        elif count in (5, 6):
            return (3, 2)
        elif count in (7, 8, 9):
            return (3, 3)
        else:
            return (1, 1)

    def initUI(self):
        self.setWindowTitle('Image comparator')
        self.resize(800, 600)
        main_layout = QVBoxLayout()

        self.grid_layout = QGridLayout()
        self.grid_layout.setContentsMargins(0, 0, 0, 0)
        self.grid_layout.setSpacing(0)
        main_layout.setContentsMargins(0, 0, 0, 0)
        main_layout.setSpacing(0)
        self.viewers = []

        rows, cols = self.get_grid_layout(len(self.image_paths))
        for i, image_path in enumerate(self.image_paths):
            viewer = ImageViewer(image_path, os.path.basename(image_path))
            self.viewers.append(viewer)
            self.grid_layout.addWidget(viewer, i // cols, i % cols)
            viewer.offset_scale_changed.connect(self.offset_scale_sync)
            viewer.position_changed.connect(self.position_sync)
            viewer.rectangle_changed.connect(self.rectangle_sync)
            viewer.delete_requested.connect(lambda v=viewer: self.delete_viewer(v))

        self.container = QWidget()
        self.container.setLayout(self.grid_layout)
        main_layout.addWidget(self.container)
        main_widget = QWidget()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)

        menubar = self.menuBar()
        file_menu = menubar.addMenu("文件(&F)")
        act_add = file_menu.addAction("添加图片(&A)...")
        act_add.triggered.connect(self.menu_add_images)
        act_clear = file_menu.addAction("清空图片(&C)")
        act_clear.triggered.connect(self.clearImages)
        act_export = file_menu.addAction("导出裁剪图(&E)...")
        act_export.triggered.connect(self.export_corp)
        act_reload = file_menu.addAction("重新加载(&R)")
        act_reload.setShortcut("F5")
        act_reload.triggered.connect(self.reloadImages)

        help_menu = menubar.addMenu("帮助(&H)")
        act_usage = help_menu.addAction("使用说明(&U)...")
        act_usage.triggered.connect(self.show_usage_dialog)

        self._toast = ToastWidget(self)

    def show_usage_dialog(self):
        prog_name = os.path.basename(sys.argv[0]) if sys.argv and sys.argv[0] else "img_cmp.py"
        text = usage(prog_name)
        dlg = UsageDialog(text, self)
        dlg.exec_()

    def show_toast(self, message: str, ms: int = 1000):
        if hasattr(self, "_toast") and self._toast is not None:
            self._toast.show_message(message, ms)

    def offset_scale_sync(self, offset: QPointF, scale: float):
        for viewer in self.viewers:
            viewer.updateGeometry(offset, scale)

    def position_sync(self, pos: QPointF):
        self.select_rect = None
        self.select_pos = pos
        for viewer in self.viewers:
            viewer.updatePosition(pos)

    def rectangle_sync(self, rect: QRectF):
        self.select_rect = rect
        
        # 找到触发信号的viewer
        sender = self.sender()
        if isinstance(sender, ImageViewer):
            self._active_viewer = sender
        
        # 更新所有viewer的显示
        for viewer in self.viewers:
            viewer.updateRectangle(rect, statistics=True)

    def align_selected_regions(self):
        """将其他图像与当前选中的图像区域对齐"""
        if not self.select_rect or not self._active_viewer:
            self.show_toast("请先选择一个区域作为对齐基准", 2000)
            return
        
        # 获取当前选区信息
        source_viewer = self._active_viewer
        target_image_size = source_viewer.image.size()
        feature_result = source_viewer.cropSubImage(self.select_rect)
        
        if feature_result is None:
            self.show_toast("无法提取选区特征，请确保选区有效", 2000)
            return
        
        # 对其他viewer进行对齐
        aligned_count = 0
        failed_viewers = []
        for viewer in self.viewers:
            if viewer != source_viewer:
                try:
                    if viewer.alignment(feature_result, source_viewer.offset, self.select_rect):
                        aligned_count += 1
                    else:
                        failed_viewers.append(viewer.name)
                except Exception as e:
                    print(f"Alignment failed for {viewer.name}: {e}")
                    failed_viewers.append(viewer.name)
        
        # 显示结果
        if aligned_count > 0:
            result_msg = f"成功对齐 {aligned_count} 个图像"
            if failed_viewers:
                result_msg += f"，{len(failed_viewers)} 个图像对齐失败"
            self.show_toast(result_msg, 3000)
        else:
            self.show_toast("没有图像成功对齐，请检查选区特征", 3000)

    def delete_viewer(self, viewer_to_delete):
        """删除指定的ImageViewer"""
        if viewer_to_delete not in self.viewers:
            return
        
        # 从布局中移除
        for i in reversed(range(self.grid_layout.count())):
            item = self.grid_layout.itemAt(i)
            if item and item.widget() == viewer_to_delete:
                self.grid_layout.removeWidget(viewer_to_delete)
                break
        
        # 断开信号连接
        viewer_to_delete.offset_scale_changed.disconnect(self.offset_scale_sync)
        viewer_to_delete.position_changed.disconnect(self.position_sync)
        viewer_to_delete.rectangle_changed.disconnect(self.rectangle_sync)
        
        # 从列表中移除
        self.viewers.remove(viewer_to_delete)
        
        # 清理对象
        viewer_to_delete.setParent(None)
        viewer_to_delete.deleteLater()
        
        # 更新图片路径列表
        self.image_paths = [v.image_path for v in self.viewers]
        
        # 重新布局
        self._reflow_grid()
        
        # 清除选区
        self.select_pos = None
        self.select_rect = None
        self._active_viewer = None
        
        # 显示提示
        self.show_toast(f"已删除图片: {viewer_to_delete.name}", 2000)

    def calculate_hist(self):
        self.show_toast("计算直方图")
        if self.select_rect is not None:
            for viewer in self.viewers:
                viewer.updateHist(self.select_rect)

    def calculate_mtf(self):
        self.show_toast("计算 MTF")
        for viewer in self.viewers:
            viewer.updateMTF(self.select_rect)

    def calculate_psnr(self):
        self.show_toast("计算 峰值信噪比")
        for viewer in self.viewers:
            viewer.updatePSNR(self.select_rect)

    def export_corp(self):
        self.show_toast("导出剪切图")
        if not self.viewers:
            return
        saveDir = QFileDialog.getExistingDirectory(self.viewers[0], "Select Export Corp Directory")
        if self.select_rect is not None and saveDir:
            for viewer in self.viewers:
                viewer.saveCorp(self.select_rect, saveDir)

    def pixel_normalization(self):
        self.show_toast("切换像素归一化\\原始数据显示")
        for viewer in self.viewers:
            viewer.showPixelNormalization()

    def toggleColorSpace(self):
        if self.color_space == 'yuv':
            self.color_space = 'rgba'
        elif self.color_space == 'rgba':
            self.color_space = 'hsv'
        else:
            self.color_space = 'yuv'
        self.show_toast(f"切换到{self.color_space}颜色空间")
        for viewer in self.viewers:
            viewer.updateColorSpace(self.color_space)

    def rotateLeft(self):
        self.show_toast("逆时针旋转90度")
        for viewer in self.viewers:
            viewer.rotateImage(clockwise=False)
        self.select_pos = None
        self.select_rect = None

    def rotateRight(self):
        self.show_toast("顺时针旋转90度")
        for viewer in self.viewers:
            viewer.rotateImage(clockwise=True)
        self.select_pos = None

    def keyPressEvent(self, event: QKeyEvent):
        key = event.key()
        if key == Qt.Key_M:
            self.calculate_mtf()
        elif key == Qt.Key_H:
            self.calculate_hist()
        elif key == Qt.Key_E:
            self.export_corp()
        elif key == Qt.Key_N:
            self.pixel_normalization()
        elif key == Qt.Key_C:
            self.toggleColorSpace()
        elif key == Qt.Key_L:
            self.rotateLeft()
        elif key == Qt.Key_R:
            self.rotateRight()
        elif key == Qt.Key_P:
            self.calculate_psnr()
        elif key == Qt.Key_A:
            self.align_selected_regions()
        elif key == Qt.Key_F5:
            self.reloadImages()

    def _reflow_grid(self):
        while self.grid_layout.count():
            item = self.grid_layout.takeAt(0)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)

        rows, cols = self.get_grid_layout(len(self.viewers))
        for i, viewer in enumerate(self.viewers):
            self.grid_layout.addWidget(viewer, i // cols, i % cols)

        self.container.setLayout(self.grid_layout)
        self.container.update()
        self.container.repaint()

    def addImage(self, image_paths):
        if not image_paths:
            return

        new_paths = [p for p in image_paths if os.path.isfile(p)]
        if not new_paths:
            return

        existed = set(self.image_paths)
        new_paths = [p for p in new_paths if p not in existed]
        if not new_paths:
            return

        all_paths = self.image_paths + new_paths
        all_paths = sorted(all_paths)
        if len(all_paths) > 9:
            all_paths = all_paths[:9]

        new_paths = [p for p in all_paths if p not in set(self.image_paths)]
        self.image_paths = all_paths

        old_by_path = {v.image_path: v for v in self.viewers}
        new_viewers = []
        for p in self.image_paths:
            if p in old_by_path:
                new_viewers.append(old_by_path[p])
            else:
                viewer = ImageViewer(p, os.path.basename(p))
                viewer.updateColorSpace(self.color_space)
                viewer.offset_scale_changed.connect(self.offset_scale_sync)
                viewer.position_changed.connect(self.position_sync)
                viewer.rectangle_changed.connect(self.rectangle_sync)
                viewer.delete_requested.connect(lambda v=viewer: self.delete_viewer(v))
                new_viewers.append(viewer)
        self.viewers = new_viewers

        self._reflow_grid()

        if self.select_pos is not None:
            for v in self.viewers:
                v.updatePosition(self.select_pos)
        elif self.select_rect is not None:
            for v in self.viewers:
                v.updateRectangle(self.select_rect, statistics=True)

    def reloadImages(self):
        self.show_toast("重新加载图片")
        for i in reversed(range(self.grid_layout.count())):
            item = self.grid_layout.takeAt(i)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)
                w.setParent(None)
        self.viewers.clear()

        rows, cols = self.get_grid_layout(len(self.image_paths))
        for i, image_path in enumerate(self.image_paths):
            viewer = ImageViewer(image_path, os.path.basename(image_path))
            viewer.updateColorSpace(self.color_space)
            viewer.offset_scale_changed.connect(self.offset_scale_sync)
            viewer.position_changed.connect(self.position_sync)
            viewer.rectangle_changed.connect(self.rectangle_sync)
            viewer.delete_requested.connect(lambda v=viewer: self.delete_viewer(v))
            self.viewers.append(viewer)
            self.grid_layout.addWidget(viewer, i // cols, i % cols)

        self.container.setLayout(self.grid_layout)
        self.container.update()
        self.container.repaint()

        self.select_pos = None
        self.select_rect = None

    def clearImages(self):
        for i in reversed(range(self.grid_layout.count())):
            item = self.grid_layout.takeAt(i)
            w = item.widget()
            if w is not None:
                self.grid_layout.removeWidget(w)
                w.setParent(None)
        self.viewers.clear()
        self.image_paths = []
        self.select_pos = None
        self.select_rect = None
        self._reflow_grid()

    def menu_add_images(self):
        paths, _ = QFileDialog.getOpenFileNames(
            self,
            "选择要添加的图片",
            "",
            "Images (*.png *.jpg *.jpeg *.bmp *.tif *.tiff *.webp);;All Files (*)"
        )
        if paths:
            self.addImage(paths)

    def dragEnterEvent(self, event: QDragEnterEvent):
        if event.mimeData().hasUrls():
            urls = event.mimeData().urls()
            local_files = [u.toLocalFile() for u in urls if u.isLocalFile()]
            if local_files:
                event.acceptProposedAction()
                return
        event.ignore()

    def dropEvent(self, event: QDropEvent):
        urls = event.mimeData().urls()
        local_files = [u.toLocalFile() for u in urls if u.isLocalFile()]
        if local_files:
            self.addImage(local_files)
            event.acceptProposedAction()
        else:
            event.ignore()


def usage(prog_name: str = "img_cmp.py") -> str:
    return f"""
用法:
python {prog_name} <image1> <image2> ... <imageN>

描述:
并排对比最多 9 张图片，支持同步缩放/拖拽、点与矩形选区、颜色信息查看、
直方图统计、MTF 计算、裁剪导出、像素归一化显示、颜色空间切换、旋转以及图像对齐。

示例:
python {prog_name} a.png b.jpg c.bmp

快捷键:
M  - 计算并显示选区的 MTF 信息
H  - 计算并显示选区的直方图
E  - 将当前选区导出另存为crop图片（弹出目录选择）
N  - 切换像素是否归一化显示(0-1)
C  - 切换颜色空间（YUV <-> RGBA <-> HSV）
P  - 计算并显示选区的 PSNR 信息
A  - 图像对齐：将其他图像与当前选中的图像区域进行像素级对齐
S  - 计算并显示选区的饱和度信息（预留）
L  - 向左旋转 90 度
R  - 向右旋转 90 度

图像对齐功能说明:
1. 在任意一个图像上使用鼠标左键拖拽选择一个区域
2. 按下快捷键 A，程序会以当前选中的图像区域为基准
3. 其他图像会自动进行SIFT特征匹配和单应性变换，实现选中区域像素级对齐，其他区可能会因变换而导致变形，而且对齐变换后的图片质量也可能发生变化，最好还是保持相同视角抓图作对比
4. 对齐完成后，所有图像会同步到相同的偏移位置

鼠标操作:
左键点击：点选，显示该点颜色信息
左键拖拽：框选矩形区域（用于图像对齐）
右键拖拽：平移
滚轮：以视图中心为基准缩放

参数:
images (位置参数)  输入的图片路径，1~9 张。

注意:
- 当前最多支持 9 张图片。
- 图像对齐功能需要选区包含足够的纹理特征，建议选择有明显边缘或纹理的区域
- 对齐效果取决于图像间的相似程度和特征点的质量
"""


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        add_help=False,
        prog="img_cmp.py"
    )
    parser.add_argument('-h', '--help', action='store_true', help='显示帮助信息并退出')
    parser.add_argument('images', nargs='*', help='图片路径，最多 9 张')

    args = parser.parse_args()

    if args.help:
        print(usage(parser.prog))
        sys.exit(0)

    image_paths = args.images[:9] if args.images else []

    invalid = [p for p in image_paths if not os.path.isfile(p)]
    if invalid:
        print("下列路径不是有效文件：")
        for p in invalid:
            print(f"  - {p}")
        print()
        print(usage(parser.prog))
        sys.exit(1)

    QCoreApplication.setAttribute(Qt.AA_EnableHighDpiScaling, True)
    QCoreApplication.setAttribute(Qt.AA_UseHighDpiPixmaps, True)

    app = QApplication(sys.argv)

    comparator = ImageComparator(image_paths)
    comparator.show()
    sys.exit(app.exec_())

